import unittest

from llvmlite import ir

from common.utils import log_and_raise
from spel4llvm.err_defs import ParseError
from spel4llvm.llvm_types import (
    LLVM_INT_TYPE,
    LLVM_FLOAT_TYPE,
    LLVM_CHAR_TYPE,
    LLVM_STR_TYPE,
    LLVM_BOOL_TYPE,
    LLVM_NULL_TYPE,
    get_llvm_list_full_type,
    is_llvm_list_type,
)


class GrammarValueType:
    VT_BOOL = 0
    VT_INT = 1
    VT_FLOAT = 2
    VT_STR = 3
    VT_CHAR = 4
    VT_NULL = 5
    VT_LIST = 6

    _TYPE_DESC = {
        VT_INT: "int",
        VT_FLOAT: "float",
        VT_STR: "str",
        VT_CHAR: "char",
        VT_BOOL: "boolean",
        VT_NULL: "null",
        VT_LIST: "list",
    }

    _BASIC_TYPE_MAPPING = {
        VT_INT: LLVM_INT_TYPE,
        VT_FLOAT: LLVM_FLOAT_TYPE,
        VT_STR: LLVM_STR_TYPE,
        VT_CHAR: LLVM_CHAR_TYPE,
        VT_BOOL: LLVM_BOOL_TYPE,
        VT_NULL: LLVM_NULL_TYPE,
    }

    def __init__(self, primary_type, elem_type=None):
        self.primary_type = primary_type
        self.elem_type = elem_type

    def __repr__(self):
        desc = self._TYPE_DESC.get(self.primary_type)
        if desc is None:
            return "unknown-type"

        if self.primary_type == self.VT_LIST:
            return "%s<%s>" % (desc, self.elem_type)
        else:
            return desc

    def __hash__(self):
        return hash((self.primary_type, self.elem_type))

    def __eq__(self, other):
        if isinstance(other, GrammarValueType):
            return (
                self.primary_type == other.primary_type
                and self.elem_type == other.elem_type
            )
        else:
            return False

    def to_llvm_type(self):
        llvm_typ = self._BASIC_TYPE_MAPPING.get(self.primary_type)
        if llvm_typ is not None:
            return llvm_typ

        if self.primary_type == self.VT_LIST:
            return get_llvm_list_full_type(self.elem_type.to_llvm_type())
        else:
            log_and_raise(TypeError, "grammar type:%s has no llvm-type" % str(self))

    # when return from llvm, we always use pointer type when collection type
    def to_llvm_return_type(self):
        llvm_typ = self.to_llvm_type()
        if is_llvm_list_type(llvm_typ):
            return ir.PointerType(llvm_typ)
        else:
            return llvm_typ

    def is_list_type(self):
        return self.primary_type == self.VT_LIST

    def is_basic_type(self):
        return self.primary_type != self.VT_LIST

    def is_numeric_type(self):
        return self.primary_type in (self.VT_INT, self.VT_FLOAT)

    @staticmethod
    def make_list_type(elem_type):
        return GrammarValueType(GrammarValueType.VT_LIST, elem_type)


# 5 basic types(NULL ignores)
# 1 complex type(list type)
# total 6 types
TP_BOOL = GrammarValueType(GrammarValueType.VT_BOOL)
TP_INT = GrammarValueType(GrammarValueType.VT_INT)
TP_FLOAT = GrammarValueType(GrammarValueType.VT_FLOAT)
TP_STR = GrammarValueType(GrammarValueType.VT_STR)
TP_CHAR = GrammarValueType(GrammarValueType.VT_CHAR)
TP_NULL = GrammarValueType(GrammarValueType.VT_NULL)


_DESC2TYPE = dict(
    (GrammarValueType._TYPE_DESC[k], k) for k in GrammarValueType._TYPE_DESC
)


def type_from_str(desc):
    vt = _DESC2TYPE.get(desc)
    if vt is None:
        log_and_raise(TypeError, "unknown type desc:%s" % desc)
    return GrammarValueType(vt)


# list-type compare without elem-type compare
def compare_type_ignore_elem_type(t1: GrammarValueType, t2: GrammarValueType):
    return t1 == t2 or (t1.is_list_type() and t2.is_list_type())


# init symbol: infer_type_from_py_value
# alloc symbol storage: py_value_to_llvm_global
# return value to jit: GrammarNode.to_llvm_return_value
# infer_type_from_py_value -> py_value_to_llvm_global -> to_llvm_return_value
def infer_type_from_py_value(py_value):
    if py_value is None:
        return None

    if isinstance(py_value, bool):
        return TP_BOOL

    if isinstance(py_value, int):
        return TP_INT

    if isinstance(py_value, float):
        return TP_FLOAT

    if isinstance(py_value, str):
        return TP_STR

    if isinstance(py_value, list):
        if py_value:
            return GrammarValueType.make_list_type(
                infer_type_from_py_value(py_value[0])
            )
        else:
            # empty list use int as elem type
            return GrammarValueType.make_list_type(TP_INT)

    log_and_raise(ParseError, "unknown python type of " + py_value)


def is_numeric_or_str(tp):
    return tp in (TP_INT, TP_FLOAT, TP_STR)


def decide_type(lhs_type: GrammarValueType, rhs_type: GrammarValueType):
    if lhs_type == rhs_type:
        return lhs_type
    elif lhs_type.is_numeric_type() and rhs_type.is_numeric_type():
        return lhs_type if lhs_type.primary_type > rhs_type.primary_type else rhs_type
    else:
        log_and_raise(
            TypeError,
            "cannot perform conversion between type:%s and %s" % (lhs_type, rhs_type),
        )


def decide_arith_type(lhs_type, rhs_type):
    if lhs_type == rhs_type:
        return lhs_type
    elif is_numeric_or_str(lhs_type) and is_numeric_or_str(rhs_type):
        return lhs_type if lhs_type.primary_type > rhs_type.primary_type else rhs_type
    else:
        log_and_raise(
            TypeError,
            "cannot perform conversion between type:%d and %d" % (lhs_type, rhs_type),
        )


class GrammarValueTypeTest(unittest.TestCase):
    def test_type_equal(self):
        self.assertTrue(
            GrammarValueType.make_list_type(TP_INT)
            == GrammarValueType.make_list_type(TP_INT)
        )
        self.assertFalse(
            GrammarValueType.make_list_type(TP_INT)
            == GrammarValueType.make_list_type(TP_STR)
        )

    def test_decide_type(self):
        self.assertEqual(decide_type(TP_INT, TP_INT), TP_INT)
        self.assertEqual(decide_type(TP_INT, TP_FLOAT), TP_FLOAT)
        self.assertRaises(TypeError, lambda: decide_type(TP_INT, TP_STR))

        list1_type = GrammarValueType.make_list_type(TP_INT)
        list2_type = GrammarValueType.make_list_type(TP_INT)
        self.assertEqual(decide_type(list1_type, list2_type), list2_type)
        self.assertEqual(decide_type(list1_type, list2_type), list1_type)

        list3_type = GrammarValueType.make_list_type(TP_STR)
        self.assertRaises(TypeError, lambda: decide_type(list1_type, list3_type))
        self.assertRaises(TypeError, lambda: decide_type(TP_STR, list3_type))

    def test_to_llvm_type(self):
        self.assertEqual(TP_STR.to_llvm_type(), LLVM_STR_TYPE)
        self.assertEqual(TP_INT.to_llvm_type(), LLVM_INT_TYPE)
        self.assertEqual(TP_FLOAT.to_llvm_type(), LLVM_FLOAT_TYPE)
        self.assertEqual(
            GrammarValueType.make_list_type(TP_INT).to_llvm_type(),
            ir.LiteralStructType([LLVM_INT_TYPE, LLVM_INT_TYPE.as_pointer()]),
        )
        self.assertEqual(
            GrammarValueType.make_list_type(TP_STR).to_llvm_type(),
            ir.LiteralStructType([LLVM_INT_TYPE, LLVM_STR_TYPE.as_pointer()]),
        )
        self.assertEqual(
            GrammarValueType.make_list_type(
                GrammarValueType.make_list_type(TP_INT)
            ).to_llvm_type(),
            ir.LiteralStructType(
                [
                    LLVM_INT_TYPE,
                    ir.LiteralStructType(
                        [LLVM_INT_TYPE, LLVM_INT_TYPE.as_pointer()]
                    ).as_pointer(),
                ]
            ),
        )

    def test_desc(self):
        self.assertEqual(str(TP_STR), "str")
        self.assertEqual(str(TP_INT), "int")

        self.assertEqual(str(GrammarValueType.make_list_type(TP_INT)), "list<int>")

        self.assertEqual(
            str(
                GrammarValueType.make_list_type(GrammarValueType.make_list_type(TP_STR))
            ),
            "list<list<str>>",
        )

    def test_infer_type_from_py_value(self):
        self.assertEqual(
            infer_type_from_py_value(True),
            TP_BOOL,
        )
        self.assertEqual(
            infer_type_from_py_value("abc"),
            TP_STR,
        )
        self.assertEqual(
            infer_type_from_py_value(3.5),
            TP_FLOAT,
        )
        self.assertEqual(
            infer_type_from_py_value(8),
            TP_INT,
        )
        self.assertEqual(
            infer_type_from_py_value(["abc", "d"]),
            GrammarValueType.make_list_type(TP_STR),
        )
        self.assertEqual(
            infer_type_from_py_value([1.5, 1]),
            GrammarValueType.make_list_type(TP_FLOAT),
        )
        self.assertEqual(
            infer_type_from_py_value([[1, 2], [3, 4]]),
            GrammarValueType.make_list_type(GrammarValueType.make_list_type(TP_INT)),
        )
        self.assertEqual(
            infer_type_from_py_value([["a", "2"], ["3", "d"]]),
            GrammarValueType.make_list_type(GrammarValueType.make_list_type(TP_STR)),
        )
        self.assertEqual(
            infer_type_from_py_value([[["a", "2"], ["3", "d"]], [["b"]]]),
            GrammarValueType.make_list_type(
                GrammarValueType.make_list_type(GrammarValueType.make_list_type(TP_STR))
            ),
        )


if __name__ == "__main__":
    unittest.main()
